# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

import math
from collections import Counter
import numpy as np
import torch
import torch.nn.functional as F

from fairseq import utils

from fairseq.criterions import FairseqCriterion, register_criterion


@register_criterion('cross_entropy_fa')
class FrequentAwareCrossEntropyCriterion(FairseqCriterion):

    def __init__(self, args, task):
        super().__init__(args, task)
        self.eos_idx = task.target_dictionary.eos()
        self.word_freq = np.zeros(len(task.target_dictionary))
        special_symbol = ['<A0>', '<V>', '<A1>', '<A2>', '<S>', '<SEP>']
        self.special_index = set([task.target_dictionary.index(sym) for sym in special_symbol])
        ents = [task.target_dictionary.index('<ent>' + str(i)) for i in range(1, 15) if task.target_dictionary.index('<ent>' + str(i)) != task.target_dictionary.unk()]
        nes = [task.target_dictionary.index('<NE>' + str(i)) for i in range(1, 15) if task.target_dictionary.index('<NE>' + str(i)) != task.target_dictionary.unk()]
        self.special_index.update(ents)
        self.special_index.update(nes)

    def forward(self, model, sample, reduce=True):
        """Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        net_output = model(**sample['net_input'])
        loss, _ = self.compute_loss(model, net_output, sample, reduce=reduce)
        sample_size = sample['target'].size(0) if self.args.sentence_avg else sample['ntokens']
        logging_output = {
            'loss': utils.item(loss.data) if reduce else loss.data,
            'ntokens': sample['ntokens'],
            'nsentences': sample['target'].size(0),
            'sample_size': sample_size,
        }
        return loss, sample_size, logging_output

    def compute_loss(self, model, net_output, sample, reduce=True):
        lprobs = model.get_normalized_probs(net_output, log_probs=True)

        # frequent loss
        preds = torch.argmax(lprobs, dim=-1)
        preds_clean = self.clean_preds(preds)
        self.update_frequency(preds_clean)
        weight = self.loss_weight()

        lprobs = lprobs.view(-1, lprobs.size(-1))
        target = model.get_targets(sample, net_output).view(-1)
        loss = F.nll_loss(lprobs, target, weight=weight, size_average=False, ignore_index=self.padding_idx,
                          reduce=reduce)
        return loss, loss

    @staticmethod
    def aggregate_logging_outputs(logging_outputs):
        """Aggregate logging outputs from data parallel training."""
        loss_sum = sum(log.get('loss', 0) for log in logging_outputs)
        ntokens = sum(log.get('ntokens', 0) for log in logging_outputs)
        nsentences = sum(log.get('nsentences', 0) for log in logging_outputs)
        sample_size = sum(log.get('sample_size', 0) for log in logging_outputs)
        agg_output = {
            'loss': loss_sum / sample_size / math.log(2),
            'ntokens': ntokens,
            'nsentences': nsentences,
            'sample_size': sample_size,
        }
        if sample_size != ntokens:
            agg_output['nll_loss'] = loss_sum / ntokens / math.log(2)
        return agg_output

    def clean_preds(self, preds):
        res = []
        preds = preds.tolist()
        for pred in preds:
            if self.eos_idx in pred:
                ind = pred.index(self.eos_idx) + 1 # end_idx included
                pred = pred[:ind]
            if len(pred) == 0:
                continue
            # if pred[0] == self.bos_idx:
            #     pred = pred[1:]
            res.append(pred)
        return res

    def reset_frequency(self):
        self.word_freq.fill(0)

    def update_frequency(self, preds):
        curr = Counter()
        for pred in preds:
            curr.update(pred)

        # self.word_freq *= self.opt['decay_factor']
        for k, v in curr.items():
            if k <= 104:
                continue
            if k == self.eos_idx:  # do not suppress END token
                continue
            if k in self.special_index:
                continue
            self.word_freq[k] += v

    def loss_weight(self):
        RF = self.word_freq / self.word_freq.sum()  # relative frequency
        a = -1 / RF.max()
        weight = a * RF + 1
        weight = weight / weight.sum() * len(weight)  # normalization
        return torch.FloatTensor(weight).cuda()